"""
SimpleOpticalFlow 测试脚本
演示光流传感器的各种功能
"""

import time
import cv2
import numpy as np
from simple_optical_flow import SimpleOpticalFlow

def test_basic_flow():
    """基础光流测试"""
    print("🚀 开始基础光流测试...")
    
    # 创建传感器
    sensor = SimpleOpticalFlow(width=640, height=480, fps=30)
    
    print("✅ 传感器已创建")
    print("⏳ 等待传感器初始化...")
    time.sleep(2)  # 等待传感器稳定
    
    try:
        print("🎯 开始获取光流数据 (按 ESC 退出)")
        
        for i in range(3000):
            # 获取光流矩阵
            flow_matrix = sensor.get_flow()
            
            if flow_matrix is not None:
                # 显示彩色光流
                sensor.show_flow_colored(step=20, scale=10, window_name="Colored Flow")
                
                # 每100帧打印一次统计信息
                if i % 100 == 0:
                    stats = sensor.calculate_flow_statistics()
                    print(f"帧 {i}: 平均光流 = {stats['mean_flow']}, "
                          f"平均幅值 = {stats['mean_magnitude']:.3f}")
            
            # 检查键盘输入
            key = cv2.waitKey(1) & 0xFF
            if key == 27:  # ESC键
                print("🛑 用户按下ESC，退出测试")
                break
                
    except KeyboardInterrupt:
        print("🛑 用户中断测试")
    
    finally:
        sensor.stop()
        cv2.destroyAllWindows()
        print("✅ 基础测试完成")

def test_all_visualizations():
    """测试所有可视化方法"""
    print("\n🎨 开始可视化测试...")
    
    sensor = SimpleOpticalFlow(width=640, height=480, fps=30)
    time.sleep(2)
    
    try:
        print("🎯 同时显示多种可视化效果 (按 ESC 退出)")
        
        for i in range(1000):
            flow_matrix = sensor.get_flow()
            
            if flow_matrix is not None:
                # 1. 黑白矢量场
                sensor.show_flow(step=20, scale=8, window_name="1. Black&White Flow")
                
                # 2. 彩色矢量场
                sensor.show_flow_colored(step=20, scale=8, window_name="2. Colored Flow")
                
                # 3. RGB编码图像
                sensor.show_flow_rgb_image(scale_factor=50, window_name="3. RGB Flow")
                
                # 每50帧显示详细信息
                if i % 50 == 0:
                    magnitude = sensor.get_flow_magnitude()
                    if magnitude is not None:
                        print(f"帧 {i}: 最大幅值 = {np.max(magnitude):.3f}, "
                              f"非零像素数 = {np.count_nonzero(magnitude > 0.1)}")
            
            key = cv2.waitKey(1) & 0xFF
            if key == 27:
                break
                
    except KeyboardInterrupt:
        print("🛑 用户中断测试")
    
    finally:
        sensor.stop()
        cv2.destroyAllWindows()
        print("✅ 可视化测试完成")

def test_statistics():
    """测试统计功能"""
    print("\n📊 开始统计功能测试...")
    
    sensor = SimpleOpticalFlow(width=640, height=480, fps=30)
    time.sleep(2)
    
    try:
        print("🎯 收集光流统计数据 (10秒)")
        
        start_time = time.time()
        frame_count = 0
        total_magnitude = 0
        max_observed_magnitude = 0
        
        while time.time() - start_time < 10:  # 运行10秒
            flow_matrix = sensor.get_flow()
            
            if flow_matrix is not None:
                frame_count += 1
                
                # 获取各种统计信息
                stats = sensor.calculate_flow_statistics()
                magnitude = sensor.get_flow_magnitude()
                direction = sensor.get_flow_direction()
                
                if magnitude is not None:
                    current_max = np.max(magnitude)
                    total_magnitude += stats['mean_magnitude']
                    max_observed_magnitude = max(max_observed_magnitude, current_max)
                    
                    # 每秒显示一次详细统计
                    if frame_count % 30 == 0:
                        print(f"📈 统计报告:")
                        print(f"   帧数: {frame_count}")
                        print(f"   平均光流: [{stats['mean_flow'][0]:.3f}, {stats['mean_flow'][1]:.3f}]")
                        print(f"   平均幅值: {stats['mean_magnitude']:.3f}")
                        print(f"   最大幅值: {stats['max_magnitude']:.3f}")
                        print(f"   幅值方差: {stats['flow_variance']:.3f}")
                        print(f"   方向范围: [{np.min(direction):.3f}, {np.max(direction):.3f}] 弧度")
                        print("   " + "-"*40)
                
                # 简单显示
                sensor.show_flow_colored(step=8, scale=5, window_name="Statistics Test")
            
            cv2.waitKey(1)
        
        # 最终统计报告
        print(f"\n📋 最终统计报告:")
        print(f"   总帧数: {frame_count}")
        print(f"   平均帧率: {frame_count/10:.1f} FPS")
        print(f"   总体平均幅值: {total_magnitude/frame_count:.3f}")
        print(f"   观察到的最大幅值: {max_observed_magnitude:.3f}")
                
    except KeyboardInterrupt:
        print("🛑 用户中断测试")
    
    finally:
        sensor.stop()
        cv2.destroyAllWindows()
        print("✅ 统计测试完成")


                

def main():
    """主测试函数"""
    print("="*60)
    print("🎮 SimpleOpticalFlow 完整测试套件")
    print("="*60)
    
    print("\n请选择测试模式:")
    print("1. 基础光流测试 (你原来的代码)")
    print("2. 可视化功能测试")
    print("3. 统计功能测试") 
    print("4. RGB编码测试")
    print("5. 光流密度配置测试")
    print("6. 运行所有测试")
    print("0. 退出")
    
    try:
        choice = input("\n请输入选择 (0-6): ").strip()
        # choice = 2
        if choice == '1':
            test_basic_flow()
        elif choice == '2':
            test_all_visualizations()
        elif choice == '3':
            test_statistics()
        elif choice == '4':
            test_rgb_encoding()

        elif choice == '0':
            print("👋 退出测试")
        else:
            print("❌ 无效选择")
            
    except KeyboardInterrupt:
        print("\n🛑 测试被中断")
    
    print("\n🎯 测试完成! 感谢使用 SimpleOpticalFlow")

if __name__ == "__main__":
    main()